import pandas as pd
from sklearn.linear_model import SGDRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score
from sklearn.preprocessing import StandardScaler  # 👈 加这句
import winsound

try:
    # 1. 读取数据
    dataset = pd.read_csv('../../dataset/比亚迪股价数据集.csv', encoding='gbk')

    # 2. 数据清洗：去除字符串中的符号，转为数值
    for col in ['开盘价', '当日最高价', '当日最低价', '交易总数', '收盘价']:
        dataset[col] = pd.to_numeric(
            dataset[col].astype(str).str.replace('[¥$,元￥]', '', regex=True),
            errors='coerce'
        )
    dataset.dropna(inplace=True)

    # 3. 构造特征
    X = dataset[['开盘价', '当日最高价', '当日最低价', '交易总数']].values
    y = dataset['收盘价'].values

    # 4. 分割数据
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42)

    # 5. ✅ 标准化特征（关键！）
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 6. 创建并训练模型
    model = SGDRegressor(max_iter=1000, tol=1e-3, random_state=42, learning_rate='adaptive', eta0=0.01)
    model.fit(X_train_scaled, y_train)

    # 7. 测试集预测
    y_pred = model.predict(X_test_scaled)
    r2 = r2_score(y_test, y_pred)
    print(f"✅ 模型 R² 分数：{r2:.4f}")

    # 8. 报警逻辑
    if r2 > 0.9:
        print("🚨 警告：模型表现非常出色")
        winsound.Beep(1000, 1000)
    else:
        print("📈 模型表现一般，但可用于预测")

    # 9. 新数据预测（必须用同样的 scaler 转换！）
    open_price = 101.00
    high_price = 105.98
    low_price = 104.10
    volume = 78420000
    new_data = [[open_price, high_price, low_price, volume]]

    # ✅ 必须用训练时的 scaler 转换新数据
    new_data_scaled = scaler.transform(new_data)
    predicted_close = model.predict(new_data_scaled)

    print(f"开盘价: {open_price}, 最高价: {high_price}, 最低价: {low_price}, 交易总数: {volume}")
    print(f"预测的收盘价: {predicted_close[0]:.2f}")

except Exception as e:
    print("❌ 发生错误：", str(e))